package dslab.client;

import dslab.ComponentFactory;
import dslab.enums.Outcome;
import dslab.util.*;

import javax.crypto.*;
import java.io.*;
import java.net.Socket;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.LinkedList;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Scanner;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class MessageClient implements IMessageClient, Runnable {

    // configuration keys that have to be present
    private static final String[] requiredConfigKeys = new String[] {
            "transfer.host", "transfer.port", "transfer.email",
            "mailbox.host", "mailbox.port", "mailbox.user", "mailbox.password"
    };

    private final String componentId;
    private final Config config;
    private final InputStream in;
    private final PrintStream out;
    private final SecretKey hmac;

    private Writer dmapWriter;
    private BufferedReader dmapReader;

    /**
     * creates a new client instance
     *
     * @param componentId the id of the component that corresponds to the Config resource
     * @param config the component config
     * @param in the input stream to read console input from
     * @param out the output stream to write console output to
     */
    public MessageClient(String componentId, Config config, InputStream in, PrintStream out) throws IOException {
        // print configuration for the benefit of the user
        System.out.printf("@%s: configuration:%n", componentId);
        for (String key : config.listKeys())
            System.out.printf("@%s: { %s: %s }%n", componentId, key, config.getString(key));

        this.componentId = componentId;
        this.config = config;
        this.in = in;
        this.out = out;

        // load hmac initially
        this.hmac = CryptoUtil.getSecretKey(Paths.get("./keys/hmac.key"));
    }

    @Override
    public void run() {
        // try-with-resources block for resource objects
        try (
            Scanner scanner = new Scanner(in);
            Socket client = new Socket(config.getString("mailbox.host"), config.getInt("mailbox.port"));
            EncryptedPrinter out = new EncryptedPrinter(new PrintWriter(client.getOutputStream(), true));
            EncryptedReader in = new EncryptedReader(new InputStreamReader(client.getInputStream()));
        ) {
            // verify whether server supports DMAP2.0
            handleInit(in);
            // handle initial AES key setup using key exchange
            performHandshake(out, in);
            // log user in
            login(out, in);

            // assign to members as we cannot change method signatures
            dmapWriter = out;
            dmapReader = in;

            String line;
            do {
                if (!client.isConnected())
                    throw new ConnectionDeadException();
                line = scanner.nextLine();
                processRequest(line);
            } while (!line.equals("shutdown"));
        } catch (NoSuchElementException ignored) {
            // thrown when the scanner gets an EOF for example, which we want to ignore gracefully
        } catch (DMAPProtocolException exception) {
            System.err.printf("@%s: a fatal error has occurred: %s%n", componentId, exception.getMessage());
        } catch (NoSuchPaddingException | NoSuchAlgorithmException exception) {
            System.err.printf("@%s: a fatal error has occurred: requested cryptographic facilities not available%n", componentId);
            exception.printStackTrace();
        } catch (BadPaddingException | IllegalBlockSizeException | InvalidAlgorithmParameterException exception) {
            System.err.printf("@%s: a fatal error has occurred: illegal usage of cryptographic functions%n", componentId);
            exception.printStackTrace();
        } catch (InvalidKeyException exception) {
            System.err.printf("@%s: a fatal error has occurred: key is invalid%n", componentId);
            exception.printStackTrace();
        } catch (IOException exception) {
            System.err.printf("@%s: a fatal error has occurred: general i/o failure%n", componentId);
            exception.printStackTrace();
        } catch (ConnectionDeadException exception) {
            System.err.printf("@%s: a fatal error has occurred: the connection has died%n", componentId);
        } finally {
            System.out.printf("sayonara%n");
            shutdown();
        }
    }

    @Override
    public void inbox() {
        // make an attempt
        try {
            String line;
            // query inbox amount
            dmapWriter.write(String.format("list%n"));
            dmapWriter.flush();
            // collect ids
            List<String> ids = new LinkedList<>();
            while (!(line = dmapReader.readLine()).startsWith("ok"))
                ids.add(expect(line, null, Pattern.compile("^(\\d+).*")).group(1));

            // show individual mails for every collected id
            for (String id : ids) {
                dmapWriter.write(String.format("show %s%n", id));
                dmapWriter.flush();

                // present the raw data to the user
                out.printf("%n%s%n", id);
                while (!(line = dmapReader.readLine()).startsWith("ok"))
                    out.println(line);
            }
        } catch (IOException | DMAPProtocolException exception) {
            System.err.printf("@%s: there was an error querying your inbox%n", componentId);
        }
    }

    @Override
    public void delete(String id) {
        // make an attempt
        try {
            // send command
            dmapWriter.write(String.format("delete %s%n", id));
            dmapWriter.flush();

            // check result
            out.println(dmapReader.readLine());
        } catch (IOException exception) {
            System.err.printf("@%s: there was an error deleting mail %s%n", componentId, id);
        }
    }

    @Override
    public void verify(String id) {
        // make an attempt
        try {
            // send command
            dmapWriter.write(String.format("show %s%n", id));
            dmapWriter.flush();

            // present the raw data to the user
            String line;
            Message message = new Message();
            line = dmapReader.readLine();
            if (line.startsWith("error")) {
                System.err.printf("error %s%n", componentId, line);
                return;
            }

            do {
                message.fill(line);
            } while (!(line = dmapReader.readLine()).startsWith("ok"));

            out.println(message);
            if (message.hash == null) {
                System.err.printf("error hash is not set for this message%n", componentId);
                return;
            }

            boolean success = CryptoUtil.verifyHash(message, hmac);
            out.printf(
                "%s message integrity %ssuccessfully verified%n",
                success ? "ok" : "error",
                success ? "" : "un-"
            );
        } catch (IOException | NoSuchAlgorithmException | InvalidKeyException exception) {
            System.err.printf("@%s: there was an error deleting mail %s%n", componentId, id);
        }
    }

    @Override
    public void msg(String to, String subject, String data) {
        // query domain
        String domain = config.getString("transfer.email").split("@")[1];

        // set message contents
        Message message = new Message();
        message.sender = config.getString("transfer.email");
        message.recipients = to;
        message.subject = subject;
        message.data = data;

        // attempt attaching a hash
        try {
            message.hash = CryptoUtil.hashMessage(message, hmac);
        } catch (NoSuchAlgorithmException | InvalidKeyException exception) {
            System.err.printf("error hash could not be attached%n");
        }

        // make DMTP client and send!
        // TODO: possibly multi-thread
        Outcome outcome = new DMTPClient(
            String.format("%s-sender-%s", componentId, domain),
            message,
            config.getString("transfer.host"),
            config.getInt("transfer.port"),
            false
        ).call();

        // notify user of what has transpired
        if (outcome != Outcome.Success)
            out.printf("error %s%n", outcome.getMessage());
        else
            out.printf("ok%n");
    }

    @Override
    public void shutdown() {
        System.out.printf("@%s: shutdown called%n", componentId);
    }

    public static void main(String[] args) throws Exception {
        IMessageClient client = ComponentFactory.createMessageClient(args[0], System.in, System.out);
        client.run();
    }

    private void handleInit(BufferedReader in) throws DMAPProtocolException,
            IllegalBlockSizeException, IOException, BadPaddingException {
        // notify user of what is happening
        System.out.printf(
            "@%s: initiating DMAP2.0 exchange%n",
            componentId
        );

        String announcement = in.readLine();
        if (!announcement.contains("ok DMAP2.0"))
            throw new DMAPProtocolException("error expected different announcement");
    }

    private void performHandshake(EncryptedPrinter out, EncryptedReader in) throws IOException,
            IllegalBlockSizeException, BadPaddingException, DMAPProtocolException, NoSuchAlgorithmException,
            InvalidAlgorithmParameterException, InvalidKeyException {
        // notify user of what is happening
        System.out.printf(
            "@%s: initiating secure handshake%n",
            componentId
        );

        // handshake step 1: C: startsecure
        out.write(String.format("startsecure%n"));
        out.flush();
        System.out.printf("@%s: handshake step 1 complete%n", componentId);

        // handshake step 2: S: ok <component-id>
        Matcher matcher;
        String step2 = in.readLine();
        matcher = expect(step2, null, Pattern.compile("ok (\\S+)"));
        Path keyPath = Paths.get(String.format("./keys/client/%s_pub.der", matcher.group(1)));
        // create cipher of public key
        Cipher cipher = CryptoUtil.createRSAEncryptCipher(CryptoUtil.getPublicKey(keyPath));
        if (cipher == null)
            throw new DMAPProtocolException("error could not create cipher with public key of server");
        System.out.printf("@%s: handshake step 2 complete: %s (key path %s)%n", componentId, step2, keyPath);

        // handshake step 3: C (RSA): ok <client-challenge> <secret-key> <iv>
        // generate components of this part of the exchange
        Pair<SecretKey, byte[]> newKeydata = CryptoUtil.generateNewKeydata();
        byte[] challenge = CryptoUtil.generateRandom(32);
        String challengeB64 = CryptoUtil.encodeB64(challenge);

        // perform encryption and encoding of binary data
        String step3 = String.format("%s%n", CryptoUtil.encryptMessage(String.format(
            "ok %s %s %s",
            challengeB64,
            CryptoUtil.encodeB64(newKeydata.key().getEncoded()),
            CryptoUtil.encodeB64(newKeydata.value())
        ), cipher));
        out.write(step3);
        out.flush();
        System.out.printf("@%s: handshake step 3 complete: %s%n", componentId, step3);

        // initialize cryptographic capabilities of streams
        out.init(newKeydata.key(), newKeydata.value());
        in.init(newKeydata.key(), newKeydata.value());

        // handshake step 4: S (AES): ok <client-challenge>
        String step4 = in.readLine();
        matcher = expect(step4, matcher, Pattern.compile("ok (\\S+)"));
        // verify if returned challenge equals sent challenge
        if (!CryptoUtil.encodeB64(challenge).equals(matcher.group(1)))
            throw new DMAPProtocolException(String.format(
                "error mailbox server has failed the challenge: expected \"%s\", got \"%s\"%n",
                CryptoUtil.encodeB64(challenge),
                matcher.group(1)
            ));
        System.out.printf("@%s: handshake step 4 complete: %s%n", componentId, step4);

        // handshake step 5: C (AES): ok
        out.write(String.format("ok%n"));
        out.flush();

        // notify user
        System.out.printf("@%s: handshake step 5 complete%n@%s: handshake successfully completed%n", componentId, componentId);
    }

    private void login(Writer out, BufferedReader in) throws IOException, DMAPProtocolException {
        // attempt login
        out.write(String.format(
            "login %s %s%n",
            config.getString("mailbox.user"),
            config.getString("mailbox.password")
        ));
        out.flush();

        // verify whether login was acknowledged
        if (!in.readLine().equals("ok"))
            throw new DMAPProtocolException("error login unsuccessful");
    }

    private Matcher expect(String line, Matcher matcher, Pattern pattern) throws DMAPProtocolException {
        // try to re-use matcher (good in a thread-safe context)
        if (matcher != null)
            matcher.reset();
        matcher = pattern.matcher(line);
        if (!matcher.find())
            throw new DMAPProtocolException(String.format(
                "error mailbox server has responded unexpectedly: expected \"%s\", got \"%s\"%n",
                pattern,
                line
            ));
        return matcher;
    }

    private void processRequest(String request) {
        String header;
        String body;

        if (request.contains(" ")) {
            // split multi-word instructions
            var headerAndBody = request.split(" ", 2);
            assert headerAndBody.length == 2;
            header = headerAndBody[0];
            body = headerAndBody[1];
        } else {
            // single word instructions have no body
            header = request;
            body = null;
        }

        // calls respective handlers and optionally does validation (to fail early)
        switch (header) {
            case "inbox":
                inbox();
                break;
            case "delete":
                if (body == null) {
                    System.err.printf("error: body required for command \"%s\"%n", header);
                    break;
                }
                delete(body);
                break;
            case "verify":
                if (body == null) {
                    System.err.printf("error: body required for command \"%s\"%n", header);
                    break;
                }
                verify(body);
                break;
            case "msg":
                if (body == null) {
                    System.err.printf("error: body required for command \"%s\"%n", header);
                    break;
                }
                Matcher matcher = Pattern.compile("^(\\S+) \"([^\"]+)\" \"([^\"]+)\"$").matcher(body);
                if (!matcher.find()) {
                    System.err.printf(
                        "cannot parse command body \"%s\" for command \"%s\"%n",
                        body,
                        header
                    );
                    break;
                }
                msg(matcher.group(1), matcher.group(2), matcher.group(3));
                break;
            case "shutdown":
                break;
            default:
                System.err.printf("command \"%s\" unknown%n", header);
                break;
        }
    }
}
